import tensorflow as tf
from module import conv3d, maxpool3d, avgpool3d


# inception模块
def inception_module(input_tensor, name_scope):
    """
    :param input_tensor: 输入的张量
    :param name_scope: 命名空间
    :return: 输出的张量
    """
    # 第一分支
    branch1 = conv3d(input_tensor, name_scope + 'branch1',
                     output_channel=64, kd=1, kh=1, kw=1, dd=1, dh=1, dw=1)
    # 第二分支
    branch2_1 = conv3d(input_tensor, name_scope + 'branch2_1',
                       output_channel=96, kd=1, kh=1, kw=1, dd=1, dh=1, dw=1)
    branch2_2 = conv3d(branch2_1, name_scope + 'branch2_2',
                       output_channel=128, kd=3, kh=3, kw=3, dd=1, dh=1, dw=1)
    # 第三分支
    branch3_1 = conv3d(input_tensor, name_scope + 'branch3_1',
                       output_channel=16, kd=1, kh=1, kw=1, dd=1, dh=1, dw=1)
    branch3_2 = conv3d(branch3_1, name_scope + 'branch3_2',
                       output_channel=32, kd=3, kh=3, kw=3, dd=1, dh=1, dw=1)
    # 第四分支
    branch4_1 = maxpool3d(input_tensor, name_scope + 'branch4_1', kd=3, kh=3, kw=3, dd=1, dh=1, dw=1)
    branch4_2 = conv3d(branch4_1, name_scope + 'branch4_2',
                       output_channel=32, kd=1, kh=1, kw=1, dd=1, dh=1, dw=1)
    # 级联各个分支
    module = tf.concat([branch1, branch2_2, branch3_2, branch4_2], 4)
    return module


# I3D，8*224*224*3
def i3d(input_tensor, num_classes, keep_prob=None):
    """
    :param input_tensor: 输入的五维张量，形如[batch_size, depth, height, width, channels]，一般为[batch_size, 8, 224, 224, 3]
    :param num_classes: 分类数
    :param keep_prob: dropout概率
    :return: 未经过softmax的一维向量，统计后的模型参数量
    """
    # 第一层，卷积层
    conv1 = conv3d(input_tensor, 'conv1', output_channel=64, kd=7, kh=7, kw=7, dd=2, dh=2, dw=2)
    pool1 = maxpool3d(conv1, 'pool1', kd=1, kh=3, kw=3, dd=1, dh=2, dw=2)
    # 第二层，卷积层
    conv2 = conv3d(pool1, 'conv2', output_channel=64, kd=1, kh=1, kw=1, dd=1, dh=1, dw=1)
    # 第三层，卷积层
    conv3 = conv3d(conv2, 'conv3', output_channel=192, kd=3, kh=3, kw=3, dd=1, dh=1, dw=1)
    pool3 = maxpool3d(conv3, 'pool3', kd=1, kh=3, kw=3, dd=1, dh=2, dw=2)
    # 第四至第十二层，inception模块
    module4 = inception_module(pool3, 'module4')
    module5 = inception_module(module4, 'module5')
    pool5 = maxpool3d(module5, 'pool5', kd=3, kh=3, kw=3, dd=2, dh=2, dw=2)
    module6 = inception_module(pool5, 'module6')
    module7 = inception_module(module6, 'module7')
    module8 = inception_module(module7, 'module8')
    module9 = inception_module(module8, 'module9')
    module10 = inception_module(module9, 'module10')
    pool10 = maxpool3d(module10, 'pool10', kd=2, kh=2, kw=2, dd=2, dh=2, dw=2)
    module11 = inception_module(pool10, 'module11')
    module12 = inception_module(module11, 'module12')
    # 保证特征图能被池化为[?, 1, 1, 1, 256]
    if module12.shape[1] > 1:
        pool12 = avgpool3d(module12, 'pool12', kd=2, kh=7, kw=7, dd=1, dh=1, dw=1, padding='VALID')
    else:
        pool12 = avgpool3d(module12, 'pool12', kd=1, kh=7, kw=7, dd=1, dh=1, dw=1, padding='VALID')
    if keep_prob is not None:
        pool12 = tf.nn.dropout(pool12, keep_prob)
    # 第十三层，卷积映射层，保证softmax的输入输出维度一致
    logit = conv3d(pool12, 'out', num_classes, kd=1, kh=1, kw=1, dd=1, dh=1, dw=1)
    # 删除指定维度
    logit = tf.squeeze(logit, [2, 3], name='squeeze')
    logit = tf.reduce_mean(logit, axis=1)
    return logit
